
#include "libutil.h"
#include "llist.h"
#include "value.h"
#include "hash_table.h"
#include "coder.h"
#include "libtimer.h"

/* Print proper usage */
static void
print_usage (const char *program)
{
  fprintf (stderr, "USAGE: %s -c | -x -m method -r runs", program);
  fprintf (stderr, " -i input_file -o output_file\n");
  fprintf (stderr, "  -c               = compress data\n");
  fprintf (stderr, "  -x               = decompress data\n");
  fprintf (stderr, "  -m [method]      = llist, htable\n");
  fprintf (stderr, "  -i [input_file]  = English text file for compress\n");
  fprintf (stderr, "                     or compressed file for decompress.\n");
  fprintf (stderr, "  -o [output_file] = Text output file for decompress or\n");
  fprintf (stderr, "                     compressed file for compress.\n");
  fprintf (stderr, "  -r [runs]        = # of runs for efficiency timing\n");
  fprintf (stderr, "  -h               = Display usage information\n");
  fprintf (stderr, "\n");
  fprintf (stderr, "EXAMPLE: %s -c -m llist -i wow.txt -o wow.pak -r 5\n",
           program);
  fprintf (stderr, "\n");
  return;
}

/* Parse arguments and check for conflicts */
void parse_cmd_args(int argc,char** argv,int* method,int* mode,
                    int* runs,char** inputfile,char** outputfile)
{
  *inputfile = *outputfile = NULL;
  *mode = UNKNOWN;*method = UNKNOWN;*runs = 1;
  int opt = GETOPT_FINISHED;
  while ((opt = getopt(argc,argv,"cxhm:i:o:r:")) != GETOPT_FINISHED) {
    switch (opt) {
      case 'i':
        *inputfile = optarg;
        break;
      case 'o':
        *outputfile = optarg;
        break;
      case 'c':
        *mode = MODE_COMPRESS;
        break;
      case 'x':
        *mode = MODE_DECOMPRESS;
        break;
      case 'm':
        if(strcmp(optarg,"llist")==0) {
          *method = METHOD_LLIST;
          fprintf(stdout,"DICTIONARY TYPE              : llist\n");
        }
        if(strcmp(optarg,"htable")==0) {
          *method = METHOD_HTABLE;
          fprintf(stdout,"DICTIONARY TYPE              : htable\n");
        }
        break;
      case 'r':
        *runs = atoi(optarg);
        break;
      case 'h':
      default:
        print_usage(argv[0]);
        exit(EXIT_FAILURE);
    }
  }
  if (*inputfile == NULL || *outputfile == NULL) {
      print_usage(argv[0]);
      exit(EXIT_FAILURE);
  }
  else if (*mode == UNKNOWN || *method == UNKNOWN) {
      print_usage(argv[0]);
      exit(EXIT_FAILURE);
  }
}

/* Initialize an empty dictionary based on the data structure
 * used.
 */
void* create_dict(int method)
{
    if(method == METHOD_LLIST) return (void*) make_empty_list();
    else return (void*) new_htable(DEF_HTABLE_SIZE);
}

/* See if the term is already in the dictionary */
value_t* check_dict_token(void* dict,char* token, int method)
{
    value_t* found;
    list_t* list = (list_t*)dict;
    htable_t* hash = (htable_t*)dict;
    if(method == METHOD_LLIST) {
       found = find_element(list,token,scmp);
    } else {
       found = htable_find_term(hash,token);
    }
    return found;
}

/* See if the id is already in the dictionary */
value_t* check_dict_id(void* dict,int id,int method)
{
    value_t* found;
    list_t* list = (list_t*)dict;
    htable_t* hash = (htable_t*)dict;
    if(method == METHOD_LLIST) {
       found = find_element(list,(void*)&id,idcmp);
    } else {
       found = htable_find_id(hash,id);
    }
    return found;
}

/* Add a new value to the dictionary by term.*/
void insert_into_dict(void* dict,value_t* new,int method)
{
    list_t* list = (list_t*)dict;
    htable_t* hash = (htable_t*)dict;
    if(method == METHOD_LLIST) {
        insert_at_foot(list,new);
    } else {
        htable_add_term(hash,new);
    }
}

/* Add a new value to the dictionary by id. */
void insert_into_dict_id(void* dict,value_t* new,int method)
{
    list_t* list = (list_t*)dict;
    htable_t* hash = (htable_t*)dict;
    if(method == METHOD_LLIST) {
        insert_at_foot(list,new);
    } else {
        htable_add_id(hash,new);
    }
}

/* Get dictionary RAM usage */
int dict_mem_usage(void* dict,int method)
{
    int size;
    list_t* list = (list_t*)dict;
    lnode_t* cur;
    htable_t* hash = (htable_t*)dict;
    if(method == METHOD_LLIST) {
        size = sizeof(list_t);
        cur = list->head;
        while(cur) {
            size += sizeof(lnode_t) + sizeof(value_t);
            size += strlen( ((value_t*)cur->data)->word);
            cur = cur->next;
        }
    } else {
        size = hash->size;
    }
    return size;
}

/* Sort link list dictionary and remap id*/
void sort_dict(void* dict,int method)
{
    list_t* list = (list_t*)dict;
    lnode_t* cur;
    unsigned int id;
    if(method == METHOD_LLIST) {
       list = sort_list(list,cntcmp); 
       /* remap ids */
       id = 0;
       cur = list->head;
       while(cur) {
            ((value_t*)cur->data)->id = id;
            id++;
            cur = cur->next;
       }
    } else {
        /* we do this in the write function */
    }
}

/* Serialize dictionary to output file */
int write_dict(void* dict,int method,FILE* out)
{
    list_t* list = (list_t*)dict;
    htable_t* hash = (htable_t*)dict;
    value_t** a;
    uint8_t token_len;
    char* word;
    int i,size;
    lnode_t* cur;
    size = 0;
    if(method == METHOD_LLIST) {
        /* dict is in freq sorted order now */
        /* write number of tokens*/
        fwrite(&list->list_len,sizeof(int),1,out);
        size += sizeof(int);
        /* iterate over tokens */
        cur = list->head;
        while(cur) {
            word = ((value_t*)cur->data)->word;
            token_len = strlen(word);
            fwrite(&token_len,sizeof(token_len),1,out);
            size += sizeof(token_len);
            fwrite(word,sizeof(char),token_len,out);
            size += sizeof(char)*token_len;
            cur = cur->next;
        }
    } else {
        a = remap_freq_htable(hash);
         /* write number of tokens*/
        fwrite(&hash->items,sizeof(int),1,out);
        size += sizeof(int);
        for(i=0;i<hash->items;i++) {
            word = (a[i])->word;
            token_len = strlen(word);
            fwrite(&token_len,sizeof(token_len),1,out);
            size += sizeof(token_len);
            fwrite(word,sizeof(char),token_len,out);
            size += sizeof(char)*token_len;
        }
        free(a);
    }
    return size;
}

/* Unserialize the dictionary */
void read_dict(void* dict,FILE* in,int method)
{
    uint8_t len;
    int n_tokens,i;
    value_t* new;

    /* read the number of tokens */
    fread(&n_tokens,sizeof(int),1,in);

    for(i=0;i<n_tokens;i++) {
        /* read token len */
        fread(&len,sizeof(len),1,in);
        new = (value_t*) safe_malloc(sizeof(value_t));
        new->word = (char*) safe_malloc(sizeof(char)*(len+1));
        fread(new->word,sizeof(char),len,in);
        /* terminate */
        new->word[len] = 0;
        new->type = isalnum(new->word[0]);
        new->count = 1;
        new->id = i;
        insert_into_dict_id(dict,new,method);
    }
}

/* Cleanup properly */
void free_dict(void* dict,int method)
{
    list_t* list;
    htable_t* hash;
    if(method == METHOD_LLIST) {
        list = (list_t*)dict;
        free_list(list,delete_value);
	free(list);
    } else {
        hash = (htable_t*)dict;
        destroy_htable(hash);
    }
}

/* Main control loop for file compression */
void compress_file(FILE* in,FILE* out,int method,int runs)
{
    char* token;
    void* dict; 
    unsigned int id = 0;
    value_t* new,*found;
    int i,type,c,n_syms,text_size,comp_size,mem_size,prelude_size;
    double *dict_const_times,*coding_times,*total_times;
    ptimer_t* timer;

    /* init timing stuff */
    dict_const_times = (double*) safe_malloc(runs*sizeof(double));
    coding_times = (double*) safe_malloc(runs*sizeof(double));
    total_times = (double*) safe_malloc(runs*sizeof(double));
    timer = new_timer();

    dict = NULL;
    prelude_size = mem_size = text_size = comp_size = 0;
    for(i=0;i<runs;i++) {
        timer_start(timer);

        /* cleanup */
        if(dict) free_dict(dict,method);
        dict = create_dict(method);

        /* calc text size */
        fseek(in,0,SEEK_END);
        text_size = ftell(in);
        fseek(in,0,SEEK_SET);
        fseek(out,0,SEEK_SET); /* reset output too */

        /* init type */
        c = fgetc(in);
        ungetc(c,in);
        type = isalnum(c);

        /* make first pass to create dict */
        n_syms = 0;
        while( (token=parse_file_char(in,&type)) != NULL ) {
            n_syms++; /* count the number of tokens */
            /* check if token is in dict */
            if( (found=check_dict_token(dict,token,method)) == NULL) {
                /* add to dict */
                new = safe_malloc(sizeof(value_t));
                new->word = token;
                new->type = type;
                new->id = id;
                new->count = 0;
                id++; 
                insert_into_dict(dict,new,method);
                found = new;
            } else {
		free(token);
	    }
            /* increase the frequency */
            found->count++;
        }

        /* calc memory usage */
        mem_size = dict_mem_usage(dict,method);


        /* sort terms by frequency so most frequent term has smallest id */
        sort_dict(dict,method);

        /* record the time it takes to construct the dict */
        dict_const_times[i] = timer_elapsed(timer);
        timer_reset(timer); /* count coding time */

        /* write dict prelude */
        prelude_size = write_dict(dict,method,out);

        /* write the number of total tokens */
        fwrite(&n_syms,sizeof(int),1,out);
        prelude_size += sizeof(int); /* part of the prelude */

        /* make second pass over text and write ids */
        rewind(in);
    
        /* init type */
        c = fgetc(in);
        ungetc(c,in);
        type = isalnum(c);

        while( (token=parse_file_char(in,&type)) != NULL ) {
            if( (found=check_dict_token(dict,token,method)) != NULL) {
                /* write id */
                bc_encode(found->id,out);
            } else {
                /* should never happen */
            }
            free(token);
        }

        coding_times[i] = timer_elapsed(timer);
        total_times[i] = coding_times[i]+ dict_const_times[i];
    }

    /* calc compressed size */
    comp_size = ftell(out); 

    /* output stuff */
    fprintf(stdout,"DICTIONARY MEMORY USAGE      : %d bytes\n",mem_size);
    fprintf(stdout,"TEXT SIZE                    : %d bytes\n",text_size);
    fprintf(stdout,"PRELUDE SIZE                 : %d bytes\n",prelude_size);
    fprintf(stdout,"COMPRESSED SIZE              : %d bytes\n",comp_size);
    /* calc compression ratio */
    fprintf(stdout,"COMPRESSION RATIO            : %.1f%%\n",
            ((float)comp_size/(float)text_size)*100);
    /* output timing */
    fprintf(stdout,"RUNS                         : %d\n",runs);
    fprintf(stdout,"DICTIONARY CONSTRUCTION TIME : %.3f +/- %.3f s\n",
            calculate_avg(dict_const_times,runs),
            calculate_stddev(dict_const_times,runs));
    fprintf(stdout,"CODING TIME                  : %.3f +/- %.3f s\n",
            calculate_avg(coding_times,runs),
            calculate_stddev(coding_times,runs));
    fprintf(stdout,"TOTAL TIME                   : %.3f +/- %.3f s\n",
            calculate_avg(total_times,runs),
            calculate_stddev(total_times,runs));

    /* cleanup */
    free_dict(dict,method);
    free(coding_times);
    free(total_times);
    free(dict_const_times);
    destroy_timer(timer);
    fclose(in);
    fclose(out);
}

/* Main control loop for file decompression */
void decompress_file(FILE* in,FILE* out,int method,int runs)
{
    void* dict;
    int n_syms,i,mem_size;
    unsigned int id;
    value_t* token;
    double *dict_const_times,*coding_times,*total_times;
    ptimer_t* timer;

    /* init timing stuff */
    dict_const_times = (double*) safe_malloc(runs*sizeof(double));
    coding_times = (double*) safe_malloc(runs*sizeof(double));
    total_times = (double*) safe_malloc(runs*sizeof(double));
    timer = new_timer();

    dict = NULL;
    mem_size = 0;
    for(i=0;i<runs;i++) {
        timer_reset(timer);

        /* seek to the front of input & output */
        fseek(in,0,SEEK_SET);
        fseek(out,0,SEEK_SET);

        /* read the dict first */
        if(dict) free_dict(dict,method);
        dict = create_dict(method);
        read_dict(dict,in,method);

        mem_size = dict_mem_usage(dict,method);

        dict_const_times[i] = timer_elapsed(timer);
        timer_reset(timer);

        /* read the number of integers we have to read */
        fread(&n_syms,sizeof(int),1,in);

        /* read encoded integers and write text */
        while(n_syms) {
            id = bc_decode(in);
            /* lookup token */
            token = check_dict_id(dict,id,method);
            if(token != NULL) {
                fprintf(out,"%s",token->word);
            } else {
                /* should never happen */
            }
            n_syms--;
        }

        coding_times[i] = timer_elapsed(timer);
        total_times[i] = coding_times[i] + dict_const_times[i];
    }

    /* output */
    fprintf(stdout,"DICTIONARY MEMORY USAGE      : %d bytes\n",mem_size);
    fprintf(stdout,"RUNS                         : %d\n",runs);
    fprintf(stdout,"DICTIONARY CONSTRUCTION TIME : %.3f +/- %.3f s\n",
            calculate_avg(dict_const_times,runs),
            calculate_stddev(dict_const_times,runs));
    fprintf(stdout,"DECODING TIME                : %.3f +/- %.3f s\n",
            calculate_avg(coding_times,runs),
            calculate_stddev(coding_times,runs));
    fprintf(stdout,"TOTAL TIME                   : %.3f +/- %.3f s\n",
            calculate_avg(total_times,runs),
            calculate_stddev(total_times,runs));

    /* cleanup */
    free(coding_times);
    free(total_times);
    free(dict_const_times);
    destroy_timer(timer);
    free_dict(dict,method);
    fclose(in);
    fclose(out);
}

int main(int argc,char** argv)
{
    FILE *fi,*fo;
    int method,mode,runs;
    char* inputfile,*outputfile;
    
    /* parse and quit there are problems */
    parse_cmd_args(argc,argv,&method,&mode,&runs,&inputfile,&outputfile);

    /* open input file */
    fi = fopen(inputfile,"r");
    fo = fopen(outputfile,"w");
    if(!fi || !fo) {
        perror("cannot open file");
        exit(EXIT_FAILURE);
    }

    if(mode == MODE_COMPRESS) {
        compress_file(fi,fo,method,runs);
    } else {
        decompress_file(fi,fo,method,runs);
    }

    return EXIT_SUCCESS;
}

